import cv2
import face_recognition
import pickle
import os

def train_face():
    # 创建数据目录
    if not os.path.exists('face_data'):
        os.makedirs('face_data')
    
    cap = cv2.VideoCapture(0)
    
    print("=== 人脸录入程序 ===")
    print("操作说明：")
    print("- 看向摄像头，按 空格键 采集人脸")
    print("- 按 Q 键退出程序")
    print("- 需要采集 5 张样本")
    print()
    
    face_encodings = []
    count = 0
    target = 5
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        # 镜像翻转
        frame = cv2.flip(frame, 1)
        
        # 检测人脸
        face_locations = face_recognition.face_locations(frame)
        
        # 画人脸框
        for (top, right, bottom, left) in face_locations:
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
        
        # 显示信息
        cv2.putText(frame, f"Samples: {count}/{target}", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(frame, "SPACE: Capture  Q: Quit", (10, 70), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        cv2.imshow('Face Training', frame)
        
        key = cv2.waitKey(1) & 0xFF
        
        # 空格键采集
        if key == ord(' ') and len(face_locations) > 0:
            encoding = face_recognition.face_encodings(frame, face_locations)[0]
            face_encodings.append(encoding)
            count += 1
            print(f"✓ 采集第 {count} 个样本")
            
            if count >= target:
                print("采集完成！")
                break
        
        # Q键退出
        elif key == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()
    
    # 保存数据
    if len(face_encodings) > 0:
        with open('face_data/user_face.pkl', 'wb') as f:
            pickle.dump(face_encodings, f)
        print(f"\n✓ 成功保存 {len(face_encodings)} 个人脸样本")
        print("✓ 数据保存至: face_data/user_face.pkl")
        print("✓ 现在可以运行 face_unlock.py 进行识别了！")
    else:
        print("\n× 未采集到人脸数据")

if __name__ == "__main__":
    train_face()